package com.utilities.IO.fileWriters;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Set;

import com.utilities.dataStructures.Range;
import com.utilities.parsers.gbkFileParser.ParseEvents;
import com.utilities.parsers.gbkFileParser.ParseGBKFileParser;
import com.utilities.parsers.gbkFileParser.ParseListener;

public class PrepareAllFeaturesFiles extends ParseListener{

	private ParseGBKFileParser fileParser = null;
	private HashMap<String,ArrayList<String>> geneMap = 
		new HashMap<String, ArrayList<String>>();
	private HashMap<String,ArrayList<String>> cdsMap = 
		new HashMap<String, ArrayList<String>>();
	private HashMap<String,ArrayList<String>> trna = 
		new HashMap<String, ArrayList<String>>();
	private HashMap<String,ArrayList<String>> ncrna = 
		new HashMap<String, ArrayList<String>>();
	private HashMap<String,ArrayList<String>> rRNA = 
		new HashMap<String, ArrayList<String>>();
	private HashMap<String,ArrayList<String>> tmRNA = 
		new HashMap<String, ArrayList<String>>();

	private String gbkFile = null;
	private String genesFile = null;
	private String cdsFile = null;
	private String ncrnafile = null;
	private String tmRnaFile = null;
	private String trnaFile = null;
	private String referenceSeqFile = null;

	private boolean collectGenes = false;
	private boolean collectCds = false;
	private boolean collectNcRNAs = false;
	private boolean colecttRNAs = false;
	private boolean collecttmRNAs = false;
	private boolean collectrRNAs = false;
	private boolean readReference = false;
	
	private FileWriter refSeqWriter = null;
	private String lastGene = null;
	private String genomeName = null;
	
	public static void main(String[] args) throws Exception
	{
		PrepareAllFeaturesFiles parser = new PrepareAllFeaturesFiles();
		parser.fileParser = new ParseGBKFileParser();
		parser.fileParser.addListener(parser);
		parser.fileParser.parse(new File(parser.gbkFile));
		if (parser.genesFile != null)
			parser.writeFiles(parser.genesFile, parser.geneMap);
		if (parser.cdsFile != null)
			parser.writeFiles(parser.cdsFile, parser.cdsMap);
		if (parser.ncrnafile != null)
			parser.writeFiles(parser.ncrnafile, parser.ncrna);
		if (parser.tmRnaFile != null)
			parser.writeFiles(parser.tmRnaFile, parser.tmRNA);
		if (parser.trnaFile != null)
			parser.writeFiles(parser.trnaFile, parser.trna);
	}

	public void run() throws IOException
	{
		fileParser = new ParseGBKFileParser();
		fileParser.addListener(this);
		fileParser.parse(new File(gbkFile));
		if (genesFile != null && collectGenes)
			writeFiles(genesFile, geneMap);
		if (cdsFile != null && collectCds)
			writeFiles(cdsFile, cdsMap);
		if (ncrnafile != null && collectNcRNAs)
			writeFiles(ncrnafile, ncrna);
		if (tmRnaFile != null && collecttmRNAs)
			writeFiles(tmRnaFile, tmRNA);
		if (trnaFile != null && collectrRNAs)
			writeFiles(trnaFile, trna);
	}

	public String getGbkFile() {
		return gbkFile;
	}

	public PrepareAllFeaturesFiles setGbkFile(String gbkFile) {
		this.gbkFile = gbkFile;
		return this;
	}

	public HashMap<String, ArrayList<String>> getGeneMap() {
		return geneMap;
	}

	public ArrayList<Range> toRange(HashMap<String, ArrayList<String>> map)
	{
		ArrayList<Range> ranges = new ArrayList<Range>();
		Set<String> keySet = map.keySet();
		for(String key : keySet)
		{
			ArrayList<String> lst = map.get(key);
			for(String entry : lst)
			{
				String pos = entry.split("\t")[0];
				int lowerLimit = Integer.parseInt(pos.split("\\.\\.")[0]);
				int upperLimit = Integer.parseInt(pos.split("\\.\\.")[1]);
				Range range = new Range(lowerLimit, upperLimit, key + "\t" + pos);
				ranges.add(range);
			}
		}
		return ranges;
	}

	public HashMap<String, ArrayList<String>> getCdsMap() {
		return cdsMap;
	}

	public HashMap<String, ArrayList<String>> getTrna() {
		return trna;
	}

	public HashMap<String, ArrayList<String>> getNcrna() {
		return ncrna;
	}

	public HashMap<String, ArrayList<String>> getrRNA() {
		return rRNA;
	}

	public HashMap<String, ArrayList<String>> getTmRNA() {
		return tmRNA;
	}

	public String getCdsFile() {
		return cdsFile;
	}

	public void setGenesFile(String genesFile) {
		this.genesFile = genesFile;
	}

	public void setNcrnafile(String ncrnafile) {
		this.ncrnafile = ncrnafile;
	}

	public void setTmRnaFile(String tmRnaFile) {
		this.tmRnaFile = tmRnaFile;
	}

	public void setTrnaFile(String trnaFile) {
		this.trnaFile = trnaFile;
	}

	public void setLastGene(String lastGene) {
		this.lastGene = lastGene;
	}

	private void writeFiles(String fileName, HashMap<String,ArrayList<String>> map) throws FileNotFoundException 
	{
		PrintWriter pw = new PrintWriter(fileName);
		Set<String> keySet = map.keySet();
		for(String key : keySet)
		{
			ArrayList<String> positionInfos = map.get(key);
			for(String pos : positionInfos)
			{
				pw.println(String.format("%s\t%s", key,pos));
			}
		}
		pw.flush();
		pw.close();
	}	

	@Override
	public void handleEvent(ParseEvents event) {
		try
		{
			switch((ParseEvents.EventTypes)event.getEventType())
			{
			case docStarted : break;
			case accessionNumber:
				if (readReference)
				{
					String str = (String)event.getSource();
					genomeName = str.trim().replaceAll("( )+", " ").split(" ")[1].trim();
				}
				break;
			case docEnded: 
			{
				if (refSeqWriter != null)
				{
					refSeqWriter.flush();
					refSeqWriter.close();
				}
				break;
			}
			case sequence:
			{
				String str = (String)event.getSource();
				String[] params = str.trim().replaceAll("( )+", " ").split(" ");
				StringBuffer strBuffer = new StringBuffer();
				System.out.println(params[0]);
				for(int i = 1; i < params.length;i++)
				{
					strBuffer.append(params[i]);
				}
				refSeqWriter.write(strBuffer.toString().trim()+"\n");
				break;
			}
			case sequenceStart: 
			{
				if (!readReference) fileParser.setStopReading(true);
				else
				{
					refSeqWriter = new FileWriter(referenceSeqFile);
					refSeqWriter.write(">" + genomeName + "\n");
				}
				break;
			}
			case rRNA:
			{
				if (colecttRNAs)
				{
					String cdsName = fileParser.getFeatureName();
					String positionInfo = fileParser.getPositionalInfo();
					// This is to check that we are not reading a gene info with no position info
					assert(positionInfo != null);
					String orientation = "+";
					if (positionInfo.contains("complement"))
					{
						orientation = "-";
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.lastIndexOf(")"));
					}
					if (positionInfo.contains("join"))
					{
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.indexOf(")")); 
						String[] subParts = positionInfo.split(",");
						String start = subParts[0].split("\\.\\.")[0];
						String end = subParts[subParts.length - 1].split("\\.\\.")[1];
						positionInfo = start+".."+end;
					}
					positionInfo = positionInfo + "\t" + orientation;
					ArrayList<String> list = null;
					// if we have already read this gene before
					if (rRNA.containsKey(cdsName))
					{
						// fetch the positions list corresponding to this geneName
						list = rRNA.get(cdsName);
					}
					else
					{
						list = new ArrayList<String>();
					}
					list.add(positionInfo);
					rRNA.put(cdsName, list);
				}
				break;
			}
			case tRNA:
			{
				if (collectrRNAs)
				{
					String cdsName = fileParser.getFeatureName();
					String positionInfo = fileParser.getPositionalInfo();
					// This is to check that we are not reading a gene info with no position info
					assert(positionInfo != null);
					String orientation = "+";
					if (positionInfo.contains("complement"))
					{
						orientation = "-";
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.lastIndexOf(")"));
					}
					if (positionInfo.contains("join"))
					{
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.indexOf(")")); 
						String[] subParts = positionInfo.split(",");
						String start = subParts[0].split("\\.\\.")[0];
						String end = subParts[subParts.length - 1].split("\\.\\.")[1];
						positionInfo = start+".."+end;
					}
					positionInfo = positionInfo + "\t" + orientation;
					ArrayList<String> list = null;
					// if we have already read this gene before
					if (trna.containsKey(cdsName))
					{
						// fetch the positions list corresponding to this geneName
						list = trna.get(cdsName);
					}
					else
					{
						list = new ArrayList<String>();
					}
					list.add(positionInfo);
					trna.put(cdsName, list);
				}
				break;
			}
			case cds: 
			{
				if (collectCds)
				{
					String cdsName = fileParser.getFeatureName();
					String positionInfo = fileParser.getPositionalInfo();
					// This is to check that we are not reading a gene info with no position info
					assert(positionInfo != null);
					String orientation = "+";
					if (positionInfo.contains("complement"))
					{
						orientation = "-";
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.lastIndexOf(")"));
					}
					if (positionInfo.contains("join"))
					{
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.indexOf(")")); 
						String[] subParts = positionInfo.split(",");
						String start = subParts[0].split("\\.\\.")[0];
						String end = subParts[subParts.length - 1].split("\\.\\.")[1];
						positionInfo = start+".."+end;
					}
					positionInfo = positionInfo + "\t" + orientation;
					ArrayList<String> list = null;
					// if we have already read this gene before
					if (cdsMap.containsKey(cdsName))
					{
						// fetch the positions list corresponding to this geneName
						list = cdsMap.get(cdsName);
					}
					else
					{
						list = new ArrayList<String>();
					}
					list.add(positionInfo);
					cdsMap.put(cdsName, list);
				}
				break;
			}
			case tmRNA:
			{
				if (collecttmRNAs)
				{
					String positionInfo = fileParser.getPositionalInfo();
					// This is to check that we are not reading a gene info with no position info
					assert(positionInfo != null);
					String orientation = "+";
					if (positionInfo.contains("complement"))
					{
						orientation = "-";
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.lastIndexOf(")"));
					}
					if (positionInfo.contains("join"))
					{
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.indexOf(")")); 
						String[] subParts = positionInfo.split(",");
						String start = subParts[0].split("\\.\\.")[0];
						String end = subParts[subParts.length - 1].split("\\.\\.")[1];
						positionInfo = start+".."+end;
					}
					positionInfo = positionInfo + "\t" + orientation;
					ArrayList<String> list = null;
					// if we have already read this gene before
					if (tmRNA.containsKey(lastGene))
					{
						// fetch the positions list corresponding to this geneName
						list = tmRNA.get(lastGene);
					}
					else
					{
						list = new ArrayList<String>();
					}
					list.add(positionInfo);
					tmRNA.put(lastGene, list);
				}
				break;
			}
			case ncRNA:
			{
				if (collectNcRNAs)
				{
					String positionInfo = fileParser.getPositionalInfo();
					// This is to check that we are not reading a gene info with no position info
					assert(positionInfo != null);
					String orientation = "+";
					if (positionInfo.contains("complement"))
					{
						orientation = "-";
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.lastIndexOf(")"));
					}
					if (positionInfo.contains("join"))
					{
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.indexOf(")")); 
						String[] subParts = positionInfo.split(",");
						String start = subParts[0].split("\\.\\.")[0];
						String end = subParts[subParts.length - 1].split("\\.\\.")[1];
						positionInfo = start+".."+end;
					}
					positionInfo = positionInfo + "\t" + orientation;
					ArrayList<String> list = null;
					// if we have already read this gene before
					if (ncrna.containsKey(lastGene))
					{
						// fetch the positions list corresponding to this geneName
						list = ncrna.get(lastGene);
					}
					else
					{
						list = new ArrayList<String>();
					}
					list.add(positionInfo);
					ncrna.put(lastGene, list);
				}
				break;
			}

			case gene:
			{
				if (collectGenes)
				{
					String geneName = fileParser.getFeatureName();
					String positionInfo = fileParser.getPositionalInfo();
					// This is to check that we are not reading a gene info with no position info
					assert(positionInfo != null);
					String orientation = "+";
					if (positionInfo.contains("complement"))
					{
						orientation = "-";
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.lastIndexOf(")"));
					}
					if (positionInfo.contains("join"))
					{
						positionInfo = positionInfo.substring(positionInfo.indexOf("(") + 1,positionInfo.indexOf(")")); 
						String[] subParts = positionInfo.split(",");
						String start = subParts[0].split("\\.\\.")[0];
						String end = subParts[subParts.length - 1].split("\\.\\.")[1];
						positionInfo = start+".."+end;
					}
					positionInfo = positionInfo + "\t" + orientation;
					ArrayList<String> list = null;
					// if we have already read this gene before
					if (geneMap.containsKey(geneName))
					{
						// fetch the positions list corresponding to this geneName
						list = geneMap.get(geneName);
					}
					else
					{
						list = new ArrayList<String>();
					}
					list.add(positionInfo);
					geneMap.put(geneName, list);
					lastGene = geneName;
				}
				break;
			}
			}
		}
		catch(Exception e)
		{
			e.printStackTrace();
			System.err.println();
			System.exit(-1);
		}
	}

	public PrepareAllFeaturesFiles setCollectGenes(boolean collectGenes) {
		this.collectGenes = collectGenes;
		return this;
	}

	public PrepareAllFeaturesFiles setCollectCds(boolean collectCds) {
		this.collectCds = collectCds;
		return this;
	}

	public PrepareAllFeaturesFiles setCollectNcRNAs(boolean collectNcRNAs) {
		this.collectNcRNAs = collectNcRNAs;
		return this;
	}

	public PrepareAllFeaturesFiles setColecttRNAs(boolean colecttRNAs) {
		this.colecttRNAs = colecttRNAs;
		return this;
	}

	public PrepareAllFeaturesFiles setCollecttmRNAs(boolean collecttmRNAs) {
		this.collecttmRNAs = collecttmRNAs;
		return this;
	}

	public PrepareAllFeaturesFiles setCollectrRNAs(boolean collectrRNAs) {
		this.collectrRNAs = collectrRNAs;
		return this;
	}

	public PrepareAllFeaturesFiles setReadReference(boolean readReference) {
		this.readReference = readReference;
		return this;
	}

	public PrepareAllFeaturesFiles setReferenceSeqFile(String referenceSeqFile) {
		this.referenceSeqFile = referenceSeqFile;
		return this;
	}
	

}
